import tiktoken
import torch
import torch.nn as nn


GPT_CONFIG_124M = {
    "vocab_size": 50257,  # 词表大小 vocabulary size
    "context_length": 1024,  # 上下文长度 context length
    "emb_dim": 768,  # 嵌入维度 embedding dimension
    "n_heads": 12,  # 注意力头数 number of attention heads
    "n_layer": 12,  # 模型层数 number of layers
    "drop_rate": 0.1,  # 丢弃概率 dropout rate
    "qkv_bias": False,  # 是否使用qkv偏置 qkv bias
}


class DummyTransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()

    def forward(self, x):
        return x


class DummyLayerNorm(nn.Module):
    def __init__(self, normalized_state, eps=1e-5):
        super().__init__()

    def forward(self, x):
        return x


class DummyGPTModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        # vocab_size 是词表大小 50257
        # emb_dim 是嵌入维度 768
        self.tok_emb = nn.Embedding(config["vocab_size"], config["emb_dim"])
        # context_length 是上下文长度
        # drop_rate 是丢弃概率
        # n_layer 是模型层数
        # n_heads 是注意力头数
        # qkv_bias 是 qkv 偏置
        self.pos_emb = nn.Embedding(config["context_length"], config["emb_dim"])
        self.drop_emb = nn.Dropout(config["drop_rate"])
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(config) for _ in range(config["n_layer"])]
        )
        self.final_norm = DummyLayerNorm(config["emb_dim"])
        self.out_head = nn.Linear(config["emb_dim"], config["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        print("input-shape: ", batch_size, seq_len)
        tok_embeds = self.tok_emb(in_idx)
        print("tok_embeds: \n", tok_embeds.shape)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        print("pos_embeds: \n", pos_embeds.shape)
        x = tok_embeds + pos_embeds
        print("x: \n", x.shape)
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits


if __name__ == "__main__":
    tokenizer = tiktoken.get_encoding("gpt2")
    batch = []
    text1 = "Every effort moves you"
    text2 = "Every day holds a"

    token_1 = tokenizer.encode(text1)
    token_2 = tokenizer.encode(text2)
    print("token_1: ", token_1)
    print("token_2: ", token_2)

    tensor_1 = torch.tensor(token_1)
    tensor_2 = torch.tensor(token_2)
    print("tensor_1: ", tensor_1)
    print("tensor_2: ", tensor_2)

    batch.append(tensor_1)
    batch.append(tensor_2)
    batch = torch.stack(batch)
    print("batch: ", batch)

    torch.manual_seed(123)
    model = DummyGPTModel(GPT_CONFIG_124M)
    logits = model(batch)
    print("logits shape: ", logits.shape)
    print("logits: ", logits)
